#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
import torch
import numpy as np
import sys
import time
import os

# ROS2相关导入
import rclpy
from rclpy.node import Node
from std_msgs.msg import Bool

os.environ['YOLO_VERBOSE'] = 'False'

class GreenBallBuzzer(Node):
    def __init__(self, model_path, camera_device=0):
        """
        初始化绿色小球蜂鸣器报警器
        """
        super().__init__('green_ball_buzzer')
        
        self.model_path = model_path
        self.camera_device = camera_device
        
        # 初始化ROS2发布器 - 只需要蜂鸣器
        print("🔄 正在初始化ROS2蜂鸣器控制...")
        try:
            self.pub_buzzer = self.create_publisher(Bool, '/Buzzer', 1)
            time.sleep(0.1)
            print("✅ ROS2蜂鸣器控制初始化成功")
        except Exception as e:
            print(f"❌ ROS2蜂鸣器控制初始化失败: {e}")
            raise
        
        # 加载YOLO模型
        print("🔄 正在加载绿色小球检测模型...")
        try:
            # 指向yolov5的目录
            self.model = torch.hub.load('/root/yolov5-7.0', 'custom', 
                                      path=model_path, source='local', force_reload=True)
            self.model.conf = 0.5
            self.model.iou = 0.45
            print(f"✅ 模型加载成功: {model_path}")
        except Exception as e:
            print(f"❌ 模型加载失败: {e}")
            raise
        
        # 初始化摄像头
        print(f"🔄 正在初始化摄像头...")
        self.cap = cv2.VideoCapture(camera_device)
        if not self.cap.isOpened():
            raise RuntimeError(f"❌ 无法打开摄像头: /dev/video{camera_device}")
        
        # 设置摄像头参数
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 854)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.cap.set(cv2.CAP_PROP_FPS, 30)
        
        self.frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print(f"✅ 摄像头初始化成功，分辨率: {self.frame_width}x{self.frame_height}")
        
        # 蜂鸣器状态
        self.buzzer_state = False
        self.detection_count = 0
        
        # 检测参数
        self.min_confidence = 0.5  # 最小置信度
        
    def publish_buzzer(self, state):
        """发布蜂鸣器状态"""
        buzzer_msg = Bool()
        buzzer_msg.data = state
        self.pub_buzzer.publish(buzzer_msg)
        
    def detect_balls(self, frame):
        """
        检测当前帧中的绿色小球
        返回是否检测到小球
        """
        results = self.model(frame)
        detections = results.pandas().xyxy[0]
        
        # 检查是否有符合条件的小球
        for _, detection in detections.iterrows():
            confidence = detection['confidence']
            if confidence > self.min_confidence:
                # 计算小球信息用于显示
                xmin, ymin = detection['xmin'], detection['ymin']
                xmax, ymax = detection['xmax'], detection['ymax']
                
                center_x = ((xmin + xmax) / 2 / self.frame_width) * 100
                center_y = ((ymin + ymax) / 2 / self.frame_height) * 100
                width = ((xmax - xmin) / self.frame_width) * 100
                height = ((ymax - ymin) / self.frame_height) * 100
                
                ball_info = {
                    'confidence': confidence,
                    'center_x': center_x,
                    'center_y': center_y,
                    'width': width,
                    'height': height
                }
                
                return True, ball_info
        
        return False, None
    
    def update_buzzer_state(self, ball_detected):
        """
        根据小球检测结果更新蜂鸣器状态
        """
        if ball_detected:
            # 检测到小球，开启蜂鸣器
            if not self.buzzer_state:
                self.publish_buzzer(True)
                self.buzzer_state = True
                print("🔊 蜂鸣器开启：检测到绿色小球")
        else:
            # 未检测到小球，关闭蜂鸣器
            if self.buzzer_state:
                self.publish_buzzer(False)
                self.buzzer_state = False
                print("🔇 蜂鸣器关闭：未检测到绿色小球")
    
    def run(self):
        """
        主运行循环
        """
        print("\n🚀 开始绿色小球蜂鸣器报警...")
        print("📝 功能说明：检测到绿色小球时蜂鸣器响铃")
        print("⌨️  按 Ctrl+C 退出\n")
        
        frame_count = 0
        last_fps_time = time.time()
        
        try:
            while rclpy.ok():
                # 读取摄像头帧
                ret, frame = self.cap.read()
                if not ret:
                    print("❌ 无法读取摄像头数据")
                    continue
                
                frame_count += 1
                
                # 检测绿色小球
                ball_detected, ball_info = self.detect_balls(frame)
                
                # 更新蜂鸣器状态
                self.update_buzzer_state(ball_detected)
                
                # 输出检测信息
                if ball_detected and ball_info:
                    self.detection_count += 1
                    print(f"[{self.detection_count:04d}] 🎯 检测到小球！置信度: {ball_info['confidence']:.1%} | "
                          f"中心: ({ball_info['center_x']:.1f}%, {ball_info['center_y']:.1f}%) | "
                          f"尺寸: {ball_info['width']:.1f}% x {ball_info['height']:.1f}% | "
                          f"蜂鸣器: {'🔊 开启' if self.buzzer_state else '🔇 关闭'}")
                else:
                    # 无检测时偶尔输出状态
                    if frame_count % 60 == 0:  # 每60帧输出一次
                        print(f"[----] ❌ 未检测到小球 | 蜂鸣器: {'🔊 开启' if self.buzzer_state else '🔇 关闭'}")
                
                # FPS统计
                current_time = time.time()
                if current_time - last_fps_time >= 5.0:
                    fps = frame_count / (current_time - last_fps_time)
                    print(f"📊 处理速度: {fps:.1f} FPS | 累计检测: {self.detection_count} 次")
                    frame_count = 0
                    last_fps_time = current_time
                
                # 处理ROS2回调
                rclpy.spin_once(self, timeout_sec=0.001)
                
        except KeyboardInterrupt:
            print("\n🛑 用户停止检测")
        except Exception as e:
            print(f"\n💥 程序运行出错: {str(e)}")
        finally:
            self.cleanup()
    
    def cleanup(self):
        """
        清理资源
        """
        print("🧹 正在清理资源...")
        
        # 关闭蜂鸣器
        for _ in range(3):
            self.publish_buzzer(False)
            time.sleep(0.05)
        
        # 释放摄像头
        if hasattr(self, 'cap') and self.cap:
            self.cap.release()
        
        cv2.destroyAllWindows()
        print(f"✅ 资源清理完成，总共检测到 {self.detection_count} 次绿色小球")

def main():
    model_path = "/root/yolov5-7.0/model/best.pt"
    
    print("🤖 绿色小球蜂鸣器报警程序 - ROS2版本")
    print("=" * 50)
    print(f"📁 模型路径: {os.path.abspath(model_path)}")
    print(f"📁 工作目录: {os.getcwd()}")
    print("🔊 报警参数:")
    print("  - 最小置信度: 0.5")
    print("  - 检测到小球时：蜂鸣器开启")
    print("  - 未检测到小球时：蜂鸣器关闭")
    print("📋 功能说明:")
    print("  - 实时检测画面中的绿色小球")
    print("  - 检测到小球时蜂鸣器响铃报警")
    print("  - 不进行任何移动控制")
    print("  - 不考虑小球在画面中的位置")
    
    # 初始化ROS2
    rclpy.init()
    
    try:
        buzzer = GreenBallBuzzer(model_path, camera_device=0)
        buzzer.run()
        
    except FileNotFoundError as e:
        print(f"📁 {str(e)}")
        print("💡 请确保模型文件 best.pt 在 model/ 目录中")
        return 1
    except Exception as e:
        print(f"💥 程序运行出错: {str(e)}")
        return 1
    finally:
        if rclpy.ok():
            rclpy.shutdown()
    
    return 0

if __name__ == "__main__":
    exit(main())